from langdev_helper.llm.lcex import llm_lcex as model

from langchain_core.tools import tool
from langgraph.graph import MessagesState, START
from langgraph.prebuilt import ToolNode
from langgraph.graph import END, StateGraph
from langgraph.checkpoint.memory import MemorySaver

@tool
def search(query: str):
    """Call to surf the web."""
    return [
        "It's sunny in San Francisco, but you better look out if you're a Gemini 😈."
    ]

tools = [search]
tool_node = ToolNode(tools)

model = model.bind_tools(tools)

def should_continue(state):
    messages = state["messages"]
    last_message = messages[-1]
    if not last_message.tool_calls:
        return "end"
    else:
        return "continue"

def call_model(state):
    messages = state["messages"]
    response = model.invoke(messages)
    return {"messages": [response]}

workflow = StateGraph(MessagesState)

workflow.add_node("agent", call_model)
workflow.add_node("action", tool_node)

workflow.add_edge(START, "agent")

workflow.add_conditional_edges(
    "agent",
    should_continue,
    {
        "continue": "action",
        "end": END,
    },
)

workflow.add_edge("action", "agent")

memory = MemorySaver()
app = workflow.compile(checkpointer=memory, interrupt_before=["action"])
print(app.get_graph().draw_mermaid())

from langchain_core.messages import HumanMessage

thread = {"configurable": {"thread_id": "3"}}
inputs = [HumanMessage(content="search for the weather in sf now")]

def main_agent():
    for event in app.stream({"messages": inputs}, thread, stream_mode="values"):
        event["messages"][-1].pretty_print()

    try:
        user_approval = input("Do you want to go to Step 3? (y/n): ")
    except:
        user_approval = "y"

    if user_approval.lower() == "y":
        # If approved, continue the graph execution
        for event in app.stream(None, thread, stream_mode="values"):
            event["messages"][-1].pretty_print()
            # print(event)
    else:
        print("Operation cancelled by user.")

    # for event in app.stream(None, thread, stream_mode="values"):
        # event["messages"][-1].pretty_print()